[None][perf] Fuse MiniMax-M3 QKV and index projection - #18205

Description

MiniMax-M3 MSA currently launches separate QKV and index projections, then
normalizes, applies RoPE, and inserts the projected K/V tensors into their
paged caches through separate producers. This adds projection and cache-write
overhead to prefill, mixed, and decode iterations and prevents the complete
projection path from integrating cleanly with piecewise CUDA graphs.

BEFORE
──────────────────────────────────────────────────────────────────────────────

                              Hidden states
                         ┌──────────┴──────────┐
                         │                     │
                         ▼                     ▼
                  Main QKV GEMM          Index-QK GEMM
                   → Q, K, V          → index-Q, index-K
                         │                     │
                         ▼                     ▼
                  Main producer         Index producer
                  • Q/K norm + RoPE     • Q/K norm + RoPE
                  • output Q            • output index-Q
                  • write current K/V   • write current index-K
                    to main KV cache      to index cache
                         │                     │
                         │                     ▼
                         │            Score visible index-K
                         │            Max-pool scores by block
                         │            Select Top-k block IDs
                         │                     │
                         └──────────┬──────────┘
                                    ▼
                         Q + selected historical
                              main K/V blocks
                                    │
                                    ▼
                          Sparse-attention kernel
                                    │
                                    ▼
                                  Output


AFTER
──────────────────────────────────────────────────────────────────────────────

                              Hidden states
                                    │
                                    ▼
                         One wider fused GEMM
                   [Q | K | V | index-Q | index-K]
                                    │
                                    ▼
                           One fused producer
                           • Q/K norm + RoPE
                           • output Q and index-Q
                           • write current K/V
                           • write current index-K
                                    │
                         ┌──────────┴──────────┐
                         │                     │
                         │                     ▼
                         │            Score visible index-K
                         │            Max-pool scores by block
                         │            Select Top-k block IDs
                         │                     │
                         └──────────┬──────────┘
                                    ▼
                         Q + selected historical
                              main K/V blocks
                                    │
                                    ▼
                          Sparse-attention kernel
                                    │
                                    ▼
                                  Output

Only the projection and producer stages are fused. Index scoring, block Top-k
selection, and the sparse-attention kernel remain separate downstream stages.
“Current K/V” is produced for the token or chunk being processed; “historical
K/V” is read from the already-populated cache using the selected block IDs.

This change adds an opt-in fused MiniMax-M3 projection path:

• packs Q, K, V, index-Q, and index-K into one quantized TP-aware projection;
• adds an SM100 producer that applies Q/K normalization and RoPE, emits compact
  Q/index-Q, and writes main K/V plus index-K directly to their paged HND FP8
  cache storage;
• supports strided cache layouts, 64-bit element offsets, CUDA Graph replay,
  and safe containment of padded or out-of-range cache slots;
• keeps the projection captured during piecewise compilation while cache
  insertion and sparse attention remain at the eager attention boundary; and
• completes the current-main one-model Eagle3 integration, including separate
  target/draft page geometry, multi-token verification, attention-DP dummy
  lifecycle, and reduced-residual hidden-state capture.

The option defaults to disabled and is accepted only by the MiniMax-M3 MSA
implementation. Fusing the projection changes weight ownership, local head
geometry, cache insertion, and the prepopulated-KV attention contract; the
producer and its Eagle3/piecewise integration therefore need to land together
to avoid an unusable intermediate configuration.

Performance

A controlled B300 aggregate TP8/EP8 attention-DP A/B used an exact 8K-input /
1K-output workload at concurrency 256 for 2,560 requests. Both arms used the
same image, model, packed projection, and runtime configuration and completed
identical input and output token counts. The isolated change extended the
fused producer from pure prefill to mixed batches and CUDA-graph decode:

This result isolates the mixed/decode extension, not the full fused path
against an unfused baseline. Its benefit is also workload-dependent: a
separate matched disaggregated GEN-only pair measured 1.32% better reciprocal
median TPOT and effectively flat total throughput (-0.31%), because it exposes
only the decode-side portion of the eliminated work.

Feature-branch origin

This current-main port consolidates the final behavior developed on the
feat/m3_with_msa side branch in four PRs:

• #16852: core fused
  projection and producer;
• #16955: mixed-batch and
  decode support;
• #17018: model integration
  and Eagle3 coverage; and
• #17227: corrected FP8
  test semantics and test registration.

These PRs are provenance, not prerequisites. This PR reconstructs their
combined result against current main.

Dependency: #16021 will provide the consolidated MiniMax-M3 MSA/Eagle3 foundation; this PR will be rebased onto it and drop overlapping runtime changes once it lands.

Test Coverage

• 52 focused B300 tests passed: 32 native/operator tests, 8
  model/Eagle3/piecewise tests, and 12 MSA backend tests.
• Full SM100 wheel build passed on B300; the exported image passed checksum,
  import, native-op, MSA, and serving-CLI validation in a fresh allocation.
• A 1,800-second B300 AgentX TP4 run with MSA FP8 caches, fused projection,
  one-model Eagle3, piecewise context graphs, and default sampling completed
  54/54 warmup and 179 profiling requests with submission_valid=true, no
  request/model/sampler errors, and no cancellations.
• Changed-file pre-commit, Python 3.12 compilation, git diff --check, and the
  regenerated LLM-args telemetry manifest passed.
• MMLU/GSM8K accuracy and Eagle3 acceptance coverage is registered for normal
  CI execution.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Dev Engineer Review

• Added an opt-in fused MiniMax-M3 Q/K/V/index-Q/index-K projection path.
• Added SM100 FP8 producer kernels for normalization, RoPE, compact Q/index-Q output, and paged HND KV-cache writes.
• Added validation for dimensions, layouts, strides, alignment, head counts, page size, device placement, and cache offsets.
• Added support for strided caches, 64-bit offsets, invalid slots, CUDA Graph replay, piecewise execution, and prepopulated KV caches.
• Added tensor-parallel-aware five-way weight sharding and loading.
• Added speculative-decoding metadata propagation and one-model Eagle3 integration.
• Updated draft-cache geometry and separate draft-cache handling.
• The fused option remains disabled by default and is restricted to MiniMax-M3 MSA.
• Configuration, operator registration, metadata handling, and supported-model documentation match the stated scope.
• No correctness, API, error-handling, or unintended-scope issue is evident from the reviewed changes.

QA Engineer Review

Added test coverage for:

• Fused projection and index-head sharding.
• Compact page-table correction and prepopulated-KV dispatch.
• Speculative-decoding scratch sizing and valid-block counts.
• Separate draft-cache geometry.
• One-engine speculative decoding, Eagle3 capture ordering, and piecewise execution.
• Five-way projection sharding and weight loading.
• FP8 horizontal production, invalid cache slots, 64-bit offsets, and KV insertion.
• MiniMax-M3 NVFP4 Eagle3 accuracy.
• Attention-DP dummy allocation rollback for target and draft KV-cache managers.

The tests/integration/test_lists/l0_b200.yml test list adds entries for the FP8 horizontal producer, indexer, and main KV insertion tests.

The tests/integration/test_lists/qa/llm_function_core.txt test list adds two MiniMax-M3 NVFP4 Eagle3 accuracy cases with attention-DP disabled and enabled.

The MSA backend, MiniMax-M3 model, FP8 producer, main KV insertion, accuracy, Eagle3, and executor test functions are not all represented by corresponding test-list entries in the provided changes.

Verdict: needs follow-up.